# 这个脚本还不能用
# 想做一个根据汉王的预测结果裁剪出单行文本的图片的功能


import os
import json
from tqdm import tqdm
import  os.path as osp
from PIL import Image
filenames = []
anns = []
pres = []
num_label = 0
num_chars = 0
dir = "E:\internship\date_test\8.25"
dst_dir = "E:\internship\date_test\8.25/crop_imghw"
os.makedirs(dst_dir,exist_ok=True)

# 保存的结果
for line in open(osp.join(dir,"hw_bo_result.txt"), "r", encoding='utf-8'):
    strs = line.split("\t")
    pres.append(json.loads(strs[1]))
# 标签
for line in open(osp.join(dir,"Label.txt.txt"), "r", encoding='utf-8'):
    strs = line.split("\t")
    filename = osp.join(dir,  strs[0])
    filenames.append(filename)
    anns.append(json.loads(strs[1]))
ftxt = open(osp.join(dst_dir, "Label.txt"), "w", encoding="utf-8")
count = 0

for i, fn in enumerate(tqdm(filenames)):
    # 文本行检测
    label_chars = []
    labe_char_boxes = []

    class_label = []
    class_pre = []
    class_pre_boxs = []
    class_label_boxs = []

    pre = pres[i]
    paraphs = []
    if "Paraph" in pre.keys():
        paraphs = pre["Paraph"]
    else:
        continue
    # 拿到所有box 预测值
    for para in paraphs:
        lines = para["line"]
        for line in lines:
            coords = line["coords"]
            labe_char_boxes.append(coords)
    if len(labe_char_boxes) == 0:
        continue

    label = anns[i]
    for j in range(len(label)):
        value = label[j]["transcription"]
        # if value == 'picture':
        if 'picture' in value:
            class_label.append(0)
        elif 'para' in value:
            class_label.append(1)
        elif 'table' in value:
            class_label.append(2)
        elif value.split("-")[0] == 'cell':
            class_label.append(3)
        else:
            label_chars.append(value)
            x1 = 1e7
            y1 = 1e7
            x2 = 0
            y2 = 0
            xyxy = label[j]["points"]
            for point in xyxy:
                if point[0] < x1:
                    x1 = point[0]
                if point[0] > x2:
                    x2 = point[0]
                if point[1] < y1:
                    y1 = point[1]
                if point[1] > y2:
                    y2 = point[1]
            labe_char_boxes.append(xyxy)
            im = Image.open(filenames[i])
            w, h = im.size
            cropim = im.crop((x1, y1-1, x2, y2+1))
            imgname = str(count)+".png"
            cropim.save(osp.join(dst_dir,imgname))
            value = value.replace("\n","")
            ftxt.write(imgname + "\t" + value +"\n")
            count = count + 1
    # 开始对比IOU
